import os
from time import sleep
import allure
from PIL import ImageGrab
import numpy as np
import cv2
import datetime
import threading

from common.log.log import log
from common.video.shots import clear_dir
from config.conf import *

flag = False  # 停止标志位


@allure.step('屏幕录制函数')
def video_record(video_file):
    """
    屏幕录制！
    :return:
    """
    try:
        p = ImageGrab.grab()  # 获得当前屏幕
        a, b = p.size  # 获得当前屏幕的大小
        fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')  # 编码格式
        video = cv2.VideoWriter(video_file, fourcc, 16, (a, b))  # 输出文件命名为test.mp4,帧率为16，可以自己设置
        while True:
            im = ImageGrab.grab()
            imm = cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)  # 转为opencv的BGR格式
            video.write(imm)
            if flag:
                log.info("录制结束！")
                log.info(f'录制的文件为：{video_file}')
                break
        video.release()
    except Exception as e:
        log.info(f'录屏过程中出现异常，为:{e}')
        raise e


@allure.step('录屏启停功能')
def stop_record_video(key=0):
    """
    键盘监听事件！！！
    尝试按squish按esc无效，改为key=1时停止录屏
    :param key:
    :return:
    """
    global flag
    # if key == keyboard.Key.esc:
    if key == 1:
        flag = True
        log.info("stop monitor！")
        return False  # 返回False，键盘监听结束！
    elif key == 0:
        flag = False
        log.info('启动录屏')


if __name__ == '__main__':
    print('开始')
    if not os.path.exists(IMG_DIR):
        os.makedirs(IMG_DIR)
    if not os.path.exists(GIF_FILE_DIR):
        os.makedirs(GIF_FILE_DIR)
    if not os.path.exists(VIDEO_DIR):
        os.makedirs(VIDEO_DIR)
    clear_dir(IMG_DIR)  # 清空临时图片目录
    clear_dir(GIF_FILE_DIR)  # 清空临时gif目录
    clear_dir(VIDEO_DIR)  # 清空临时录屏目录
    name = datetime.datetime.now().strftime('%Y-%m-%d %H-%M-%S')  # 当前的时间
    file = VIDEO_DIR + f'\\{name}.mp4'
    print(f'录屏文件为{file}')
    th = threading.Thread(target=video_record, args=(file,), )
    th.start()
    print('开始录屏')
    print('录制中。。。。')
    sleep(5)
    stop_record_video(1)  # 录制5秒后停止
